from torch.utils.data import DataLoader

import torchvision
from torchvision import transforms
from torchvision.datasets import CIFAR10

#import torch
#import torch.nn as nn
#import torch.nn.functional as f

#import torch.optim as optim
#import time
#import os

import matplotlib.pyplot as plt
import numpy as np

import config as cfg

# RandomHorizontalFlip和RandomGrayscale用来做数据增强，防止训练出现过拟合；
# 通常在小型数据集上，通过随机翻转图片，随机调整图片的亮度，来达到增加训练时数据集的容量。
# Normalize 给定均值：(R,G,B) 方差：（R，G，B），将会把Tensor正则化。
transform_random = transforms.Compose([
    transforms.RandomHorizontalFlip(),
    transforms.RandomGrayscale(),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

transform_norm = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

# root是数据集下载保存的位置
# train_data = CIFAR10(root='./dataset', train=True, transform=transforms.ToTensor(), download=True)
# test_data = CIFAR10(root='./dataset', train=False, transform=transforms.ToTensor(), download=True)
# test_loader = DataLoader(test_data, batch_size=64)
# train_loader = DataLoader(train_data, batch_size=64)

classes_CIFAR10 = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')


def transform():
    return transform_norm

def classes():
    return classes_CIFAR10

def get_loader(train: bool):
    train_data = CIFAR10(root=cfg.file_dataset, train=train, transform=transform(), download=True)
    print(f'数据集长度:{len(train_data)} train:{train}')
    return DataLoader(train_data, batch_size=cfg.batch_size, shuffle=True, num_workers=cfg.num_thread)


def imshow(img):
    img = img / 2 + 0.5  # unnormalize
    npimg = img.numpy()
    plt.imshow(np.transpose(npimg, (1, 2, 0)))
    plt.show()


def get_some_images(train_loader):
    # get some random training images
    dataiter = iter(train_loader)
    images, labels = dataiter.next()
    imshow(torchvision.utils.make_grid(images))
    # print labels
    print(' '.join('%5s' % classes[labels[j]] for j in range(4)))
